from neo4j import GraphDatabase
import csv

uri = "bolt://localhost:7687"
username = "neo4j"
password = "wxt200413"

def import_data(file_path):
    try:
        driver = GraphDatabase.driver(uri, auth=(username, password))
        with driver.session() as session, open(file_path, 'r', encoding='utf-8') as file:
            csv_reader = csv.DictReader(file)
            for row in csv_reader:
                # 提取技术栈字符串，并拆分成多个名词
                tech_stack_str = row['技术栈']
                tech_stack_list = [word.strip() for word in tech_stack_str.split(',')]

                # 去除重复的名词
                tech_stack_list = list(set(tech_stack_list))

                query = """
                    CREATE (c:公司 {名称: $company})
                    CREATE (a:地址 {名称: $address})
                    CREATE (c)-[:地址]->(a)
                    CREATE (t:职位 {id: $id, 名称: $title, 学历要求: $education, 薪资: $salary, 链接: $link}) 
                    CREATE (a)-[:招聘岗位]->(t)
                    CREATE (e:要求 {name: $description})
                    CREATE (t)-[:招聘要求]->(e)
                    MERGE (job:工作 {name: '计算机工作招聘'})  
                    CREATE (c)-[:PART_OF]->(job)  
                    """
                params = {
                    "id": row['ID'],
                    "company": row['公司名称'],
                    "address": row['address'],
                    "salary": row['salary'],
                    "title": row['岗位'],
                    "link": row['link'],
                    "education": row['education'],
                    "description": row['description']
                }
                session.run(query, params)
                # 创建每个名词节点，并用 '提取' 关系连接到要求节点上
                for word in tech_stack_list:
                    query = """
                    MATCH (e:要求 {name: $description})
                    MERGE (s:技术栈 {name: $name}) 
                    CREATE (e)-[:提炼]->(s)   
                    """
                    params = {
                        "name": word,
                        "description": row['description']
                    }
                    session.run(query, params)
        print("Data imported successfully.")
    except Exception as e:
        print(f"An error occurred: {e}")
    finally:
        driver.close()

import_data('D:/IDEProjects/Job_Recommend/Files/Prepare/招聘信息.csv')

